import os.path
import sys
from Bio import SeqIO
from Bio import SeqUtils
from Bio import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Align.Applications import ClustalwCommandline

def is_fasta(filename):
    fasta_ext = [".fasta", ".fa"]
    return os.path.splitext(filename)[1] in fasta_ext

def is_gbk(filename):
    gbk_ext = [".gbk", ".gb"]
    return os.path.splitext(filename)[1] in gbk_ext

def read_from_file(filename):
    records = []
    with open(filename) as handle:
        if is_fasta(filename):
            records += list(SeqIO.parse(handle, "fasta"))
        else:
            records += list(SeqIO.parse(handle, "genbank"))
    return records

def get_gc_content(seq):
    return SeqUtils.GC(seq)

def transcribe_sequence(seq):
    return Seq.transcribe(seq)

def translate_sequence_to_seqrecord(seq, id):
    return SeqRecord(Seq.translate(seq), id=id)

def multiple_alignment(infile, outfile_prefix):
    cline = ClustalwCommandline("clustalw2", infile = infile, matrix = "BLOSUM", outfile = outfile_prefix + ".aln")
    stdout, stderr = cline()
    aln_records = SeqIO.parse(open(outfile_prefix + ".aln"), "clustal")
    SeqIO.write(aln_records, open(outfile_prefix + ".phy", "w"), "phylip")
    return stdout

def main(filenames):
    if not filenames:
        sys.stderr.write("No input file given\n")
        return -1

    for f in filenames:
        if not is_fasta(f) and not is_gbk(f):
            print >> sys.stderr, "One of input files has wrong format: " + f
            return -2

    seq_records = {}

    print "Reading files"
    for f in filenames:
        seq_records[f] = read_from_file(f)

    print "\nCalculating GC:"
    for (filename, records) in seq_records.items():
        print filename + ":"
        for record in records:
            print "\t" + record.id + ": " + str(get_gc_content(record.seq))

    print "\nTranscribing:"
    for (filename, records) in seq_records.items():
        print filename + ":"
        for record in records:
            print "\t" + record.id + ": "
            print transcribe_sequence(record.seq)

    print "\nTranslating:"
    protein_recs = []
    for (filename, records) in seq_records.items():
        print filename + ":"
        for record in records:
            print "\t" + record.id + ": "
            protein_recs.append(translate_sequence_to_seqrecord(record.seq, str(len(protein_recs)) + "_" + record.id))
            print protein_recs[-1].seq
    SeqIO.write(protein_recs, open("protein.fasta", "w"), "fasta")

    print "\nMultiple alignment:"
    print multiple_alignment("protein.fasta", "protein")
    return 0

if __name__ == "__main__":
    filenames = sys.argv[1:]
    
    exit_code = main(filenames)
    sys.exit(exit_code)
